{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.matlib\n",
    "import torch\n",
    "import math\n",
    "import pdb\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e_path = '/home/sci/amanpreet/Documents/DT_KG/Codes/Registered/Testing/ellipse.tif'\n",
    "c_path = '/home/sci/amanpreet/Documents/DT_KG/Codes/Registered/Testing/circle.tif'\n",
    "\n",
    "ellipse = np.asarray(Image.open(e_path))\n",
    "circle = np.asarray(Image.open(c_path))\n",
    "\n",
    "print(circle.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, (ax1, ax2) = plt.subplots(1, 2)\n",
    "ax1.imshow(circle)\n",
    "ax2.imshow(ellipse)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.max(circle))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.array([1 ,0 ,-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grad(f):\n",
    "    f = f.view(1,1,f.shape[0],f.shape[1])\n",
    "    x_filter = torch.from_numpy(np.reshape(np.array([[0,0,0],[1 ,0 ,-1],[0,0,0]]),[1,1,3,3])).type(torch.FloatTensor).cuda()\n",
    "    y_filter = torch.from_numpy(np.reshape(np.array([[0,1,0],[0 ,0 ,0],[0,-1,0]]),[1,1,3,3])).type(torch.FloatTensor).cuda()\n",
    "    ux = F.conv2d(f,x_filter,padding=1)\n",
    "    uy = F.conv2d(f,y_filter,padding=1)\n",
    "    return ux[0,0,:,:], uy[0,0,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp2(f,hy,hx):\n",
    "    f = f.view(1,1,f.shape[0],f.shape[1])\n",
    "       \n",
    "    grid = torch.stack([hy,hx], dim=2)\n",
    "    grid = grid.view(1,grid.shape[0],grid.shape[1],2).type(torch.cuda.FloatTensor)\n",
    "\n",
    "    new_im = F.grid_sample(f,grid)\n",
    "    return new_im[0,0,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elastic_deformation(i1, i2, sigma, l, n_iter, step_size):\n",
    "    lookup_field_y = np.matlib.repmat(np.arange(i1.shape[1]),i1.shape[0],1)\n",
    "    lookup_field_x = np.transpose(np.matlib.repmat(np.arange(i1.shape[0]),i1.shape[1],1))\n",
    "    \n",
    "    def_field_x = np.zeros(lookup_field_x.shape)\n",
    "    def_field_y = np.zeros(lookup_field_y.shape)\n",
    "    \n",
    "    [r,c] = i1.shape\n",
    "    L = np.zeros(lookup_field_x.shape)\n",
    "    \n",
    "    #Set up the Laplacian for onv Laplacian operation\n",
    "    for u in range(r):\n",
    "        for v in range(c):\n",
    "            L[u][v] = -2.0*(math.cos(2*math.pi*u/r)+math.cos(2*math.pi*v/c))+4.0\n",
    "            \n",
    "    #Send everything to torch Tensors and then to GPU\n",
    "    \n",
    "    i1 = torch.from_numpy(i1/255.0).type(torch.cuda.FloatTensor)\n",
    "    i2 = torch.from_numpy(i2/255.0).type(torch.cuda.FloatTensor)\n",
    "    lookup_field_x = torch.from_numpy(lookup_field_x).type(torch.cuda.FloatTensor)\n",
    "    lookup_field_y = torch.from_numpy(lookup_field_y).type(torch.cuda.FloatTensor)\n",
    "    \n",
    "    #Different identity lookup field in pytorch\n",
    "    \n",
    "    lookup_field_x = (lookup_field_x-lookup_field_x.mean())/lookup_field_x.mean() \n",
    "    lookup_field_y = (lookup_field_y-lookup_field_y.mean())/lookup_field_y.mean()\n",
    "    \n",
    "    def_field_x = torch.from_numpy(def_field_x).type(torch.cuda.FloatTensor)\n",
    "    def_field_y = torch.from_numpy(def_field_y).type(torch.cuda.FloatTensor)\n",
    "    L = torch.from_numpy(L).type(torch.cuda.FloatTensor)\n",
    "    \n",
    "    # Compute gradient of i2\n",
    "    \n",
    "    [i2_x, i2_y] = grad(i2)\n",
    "    \n",
    "    #Start the loop\n",
    "    \n",
    "    for iter in range(n_iter):\n",
    "        hx = lookup_field_x + def_field_x\n",
    "        hy = lookup_field_y + def_field_y\n",
    "\n",
    "        i2_u = interp2(i2,hy,hx);\n",
    "\n",
    "        diff = i2_u-i1    \n",
    "\n",
    "        #Compute the gradient at the deformation\n",
    "\n",
    "        grad_i2_x = interp2(i2_x,hy,hx)\n",
    "        grad_i2_y = interp2(i2_y,hy,hx)\n",
    "\n",
    "        der_x = diff * grad_i2_y\n",
    "        der_y = diff * grad_i2_x\n",
    "\n",
    "        # Now take the der to fourier space and apply inverse laplacian\n",
    "\n",
    "        der_x_f = torch.rfft(der_x,2,onesided=False)       \n",
    "        der_x_f[:,:,0] = der_x_f[:,:,0]/(L+0.0001)\n",
    "        der_x_f[:,:,1] = der_x_f[:,:,1]/(L+0.0001)\n",
    "\n",
    "        der_y_f = torch.rfft(der_y,2,onesided=False)       \n",
    "        der_y_f[:,:,0] = der_y_f[:,:,0]/(L+0.0001)\n",
    "        der_y_f[:,:,1] = der_y_f[:,:,1]/(L+0.0001)\n",
    "\n",
    "        # Now take it back to spatial domain.\n",
    "\n",
    "        der_x_f = torch.irfft(der_x_f,2,onesided=False)\n",
    "        der_y_f = torch.irfft(der_y_f,2,onesided=False)    \n",
    "\n",
    "        def_field_x += (step_size/(sigma**2 * l))*der_x_f\n",
    "        def_field_y += (step_size/(sigma**2 * l))*der_y_f\n",
    "        \n",
    "        if iter%10 == 0:\n",
    "            f, (ax1, ax2, ax3) = plt.subplots(1, 3)\n",
    "            ax2.imshow(i2_u.cpu().numpy(), aspect=\"auto\")\n",
    "            ax2.set_xticks([])\n",
    "            ax2.set_yticks([])\n",
    "            ax3.imshow(diff.cpu().numpy(), aspect=\"auto\")\n",
    "            ax3.set_xticks([])\n",
    "            ax3.set_yticks([])\n",
    "            ax1.set_xlim(-1,1)\n",
    "            ax1.set_ylim(-1,1)\n",
    "            ax1.plot(hy.cpu().numpy()[0::10,0::10],hx.cpu().numpy()[0::10,0::10],'b')\n",
    "            ax1.plot(np.transpose(hy.cpu().numpy()[0::10,0::10]),np.transpose(hx.cpu().numpy()[0::10,0::10]),'b')\n",
    "            ax1.set_xticks([])\n",
    "            ax1.set_yticks([])\n",
    "            ax1.invert_yaxis()\n",
    "            f.set_size_inches(100, 10)\n",
    "            plt.show()\n",
    "            print(\"Difference is\", np.sum(diff.cpu().numpy()))\n",
    "    \n",
    "    hx = lookup_field_x + def_field_x\n",
    "    hy = lookup_field_y + def_field_y\n",
    "    return [interp2(i2,hy,hx).cpu().numpy(), hx.cpu().numpy(), hy.cpu().numpy()]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[a,b,c] = elastic_deformation(circle,ellipse,500,1,100,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im1_path = '/home/sci/amanpreet/Documents/DT_KG/Codes/Registered/828_24.tif'\n",
    "im2_path = '/home/sci/amanpreet/Documents/DT_KG/Codes/Registered/829_22.tif'\n",
    "im1 = np.asarray(Image.open(im1_path))\n",
    "im2 = np.asarray(Image.open(im2_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[reg_img, hx, hy] = elastic_deformation(im1,im2,500,1,200,200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diffeomorphic_deformation(i1, i2, sigma, l, n_iter, step_size):\n",
    "    lookup_field_y = np.matlib.repmat(np.arange(i1.shape[1]),i1.shape[0],1)\n",
    "    lookup_field_x = np.transpose(np.matlib.repmat(np.arange(i1.shape[0]),i1.shape[1],1))\n",
    "    \n",
    "    def_field_x = np.zeros(lookup_field_x.shape)\n",
    "    def_field_y = np.zeros(lookup_field_y.shape)\n",
    "    \n",
    "    [r,c] = i1.shape\n",
    "    L = np.zeros(lookup_field_x.shape)\n",
    "    \n",
    "    #Set up the Laplacian for onv Laplacian operation\n",
    "    for u in range(r):\n",
    "        for v in range(c):\n",
    "            L[u][v] = -2.0*(math.cos(2*math.pi*u/r)+math.cos(2*math.pi*v/c))+4.0\n",
    "            \n",
    "    #Send everything to torch Tensors and then to GPU\n",
    "    \n",
    "    i1 = torch.from_numpy(i1/255.0).type(torch.cuda.FloatTensor)\n",
    "    i2 = torch.from_numpy(i2/255.0).type(torch.cuda.FloatTensor)\n",
    "    lookup_field_x = torch.from_numpy(lookup_field_x).type(torch.cuda.FloatTensor)\n",
    "    lookup_field_y = torch.from_numpy(lookup_field_y).type(torch.cuda.FloatTensor)\n",
    "    \n",
    "    #Different identity lookup field in pytorch\n",
    "    \n",
    "    lookup_field_x = (lookup_field_x-lookup_field_x.mean())/lookup_field_x.mean() \n",
    "    lookup_field_y = (lookup_field_y-lookup_field_y.mean())/lookup_field_y.mean()\n",
    "    \n",
    "    def_field_x = torch.from_numpy(def_field_x).type(torch.cuda.FloatTensor)\n",
    "    def_field_y = torch.from_numpy(def_field_y).type(torch.cuda.FloatTensor)\n",
    "    L = torch.from_numpy(L).type(torch.cuda.FloatTensor)\n",
    "    \n",
    "    # Compute gradient of i2\n",
    "    \n",
    "    [i2_x, i2_y] = grad(i2)\n",
    "    \n",
    "    #Start the loop\n",
    "    \n",
    "    for iter in range(n_iter):\n",
    "        hx = lookup_field_x + def_field_x\n",
    "        hy = lookup_field_y + def_field_y\n",
    "\n",
    "        i2_u = interp2(i2,hy,hx);\n",
    "\n",
    "        diff = i2_u-i1    \n",
    "\n",
    "        #Compute the gradient at the deformation\n",
    "\n",
    "        grad_i2_x = interp2(i2_x,hy,hx)\n",
    "        grad_i2_y = interp2(i2_y,hy,hx)\n",
    "\n",
    "        der_x = diff * grad_i2_y\n",
    "        der_y = diff * grad_i2_x\n",
    "\n",
    "        # Now take the der to fourier space and apply inverse laplacian\n",
    "\n",
    "        der_x_f = torch.rfft(der_x,2,onesided=False)       \n",
    "        der_x_f[:,:,0] = der_x_f[:,:,0]/(L+0.0001)\n",
    "        der_x_f[:,:,1] = der_x_f[:,:,1]/(L+0.0001)\n",
    "\n",
    "        der_y_f = torch.rfft(der_y,2,onesided=False)       \n",
    "        der_y_f[:,:,0] = der_y_f[:,:,0]/(L+0.0001)\n",
    "        der_y_f[:,:,1] = der_y_f[:,:,1]/(L+0.0001)\n",
    "\n",
    "        # Now take it back to spatial domain.\n",
    "\n",
    "        der_x_f = torch.irfft(der_x_f,2,onesided=False)\n",
    "        der_y_f = torch.irfft(der_y_f,2,onesided=False)    \n",
    "\n",
    "        velocity_f_x = (step_size/(sigma**2 * l))*der_x_f\n",
    "        velocity_f_y = (step_size/(sigma**2 * l))*der_y_f\n",
    "        \n",
    "        cx = lookup_field_x + velocity_f_x\n",
    "        cy = lookup_field_y + velocity_f_y\n",
    "        \n",
    "        cx[cx > 1] = 1\n",
    "        cx[cx < -1] = -1\n",
    "        \n",
    "        cy[cy > 1] = 1\n",
    "        cy[cy < -1] = -1\n",
    "        \n",
    "        def_field_x = interp2(hx, cy, cx) - lookup_field_x\n",
    "        def_field_y = interp2(hy, cy, cx) - lookup_field_y        \n",
    "        \n",
    "        if iter%10 == 0:\n",
    "            f, (ax1, ax2, ax3) = plt.subplots(1, 3)\n",
    "            ax2.imshow(i2_u.cpu().numpy(), aspect=\"auto\")\n",
    "            ax2.set_xticks([])\n",
    "            ax2.set_yticks([])\n",
    "            ax3.imshow(diff.cpu().numpy(), aspect=\"auto\")\n",
    "            ax3.set_xticks([])\n",
    "            ax3.set_yticks([])\n",
    "            ax1.set_xlim(-1,1)\n",
    "            ax1.set_ylim(-1,1)\n",
    "            ax1.plot(hy.cpu().numpy()[0::10,0::10],hx.cpu().numpy()[0::10,0::10],'b')\n",
    "            ax1.plot(np.transpose(hy.cpu().numpy()[0::10,0::10]),np.transpose(hx.cpu().numpy()[0::10,0::10]),'b')\n",
    "            ax1.set_xticks([])\n",
    "            ax1.set_yticks([])\n",
    "            ax1.invert_yaxis()\n",
    "            f.set_size_inches(100, 10)\n",
    "            plt.show()\n",
    "            print(\"Difference is\", np.sum(diff.cpu().numpy()))\n",
    "    \n",
    "    hx = lookup_field_x + def_field_x\n",
    "    hy = lookup_field_y + def_field_y\n",
    "    return [interp2(i2,hy,hx).cpu().numpy(), hx.cpu().numpy(), hy.cpu().numpy()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[a,b,c] = diffeomorphic_deformation(circle,ellipse,1000,1,100,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
